import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

# 创建SummaryWriter实例，指定日志输出目录
writer = SummaryWriter('logs')

# 数据预处理：将图像调整为32x32并转换为张量
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor()
])

# 加载MNIST训练集
train_dataset = datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transform
)

# 创建数据加载器，批量大小设为100
train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=100,
    shuffle=False
)


step = 0
for images, labels in train_loader:
    # 创建图像网格（10x10）
    img_grid = torchvision.utils.make_grid(images, nrow=10)
    img_grid = torch.reshape(img_grid, (1, 1, 342*3, 342))
    writer.add_images('images', img_grid, step)
    step += 1


writer.close()

print("数据已写入TensorBoard日志目录，可使用以下命令启动TensorBoard：")
print("tensorboard --logdir=test/logs")
